import argparse
import gzip
import os
import pickle

import numpy as np


def get_data():
    # 文件获取
    train_image = r"/workspace/kubeflow-mnist/datasets/fashion-mnist/train-images-idx3-ubyte.gz"
    test_image = r"/workspace/kubeflow-mnist/datasets/fashion-mnist/t10k-images-idx3-ubyte.gz"
    train_label = r"/workspace/kubeflow-mnist/datasets/fashion-mnist/train-labels-idx1-ubyte.gz"
    test_label = r"/workspace/kubeflow-mnist/datasets/fashion-mnist/t10k-labels-idx1-ubyte.gz"
    paths = [train_label, train_image, test_label, test_image]

    with gzip.open(paths[0], 'rb') as lbpath:
        y_train = np.frombuffer(lbpath.read(), np.uint8, offset=8)

    with gzip.open(paths[1], 'rb') as imgpath:
        x_train = np.frombuffer(
            imgpath.read(), np.uint8, offset=16).reshape(len(y_train), 28, 28)

    with gzip.open(paths[2], 'rb') as lbpath:
        y_test = np.frombuffer(lbpath.read(), np.uint8, offset=8)

    with gzip.open(paths[3], 'rb') as imgpath:
        x_test = np.frombuffer(
            imgpath.read(), np.uint8, offset=16).reshape(len(y_test), 28, 28)

    return (x_train, y_train), (x_test, y_test)


def preprocess(data_dir: str):
    print("...开始获取数据集......")
    (train_images, train_labels), (test_images, test_labels) = get_data()
    print("...数据集下载完毕......")
    train_images = train_images / 255.0
    test_images = test_images / 255.0

    print("...开始创建文件夹......")
    os.makedirs(data_dir, exist_ok=True)

    print("...开始创建训练集X文件......")
    with open(os.path.join(data_dir, 'train_images.pickle'), 'wb') as f:
        pickle.dump(train_images, f)

    print("...开始创建训练集Y文件......")
    with open(os.path.join(data_dir, 'train_labels.pickle'), 'wb') as f:
        pickle.dump(train_labels, f)

    print("...开始创建测试集X文件......")
    with open(os.path.join(data_dir, 'test_images.pickle'), 'wb') as f:
        pickle.dump(test_images, f)

    print("...开始创建测试集Y文件......")
    with open(os.path.join(data_dir, 'test_labels.pickle'), 'wb') as f:
        pickle.dump(test_labels, f)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Kubeflow MNIST training script')
    parser.add_argument('--data_dir', help='path to images and labels.')
    args = parser.parse_args()

    preprocess(data_dir=args.data_dir)
